# =============================================================================
# Structural_3lights.R — MLE strategy estimation with stratEst
#                         for the Shaded Lights Experiment (3 Lights)
#
# Strategies included: single-variable and two-variable rules only.
# Three-variable rules are excluded.
#
# OBSERVABLE STATES (8 total — one per unique (Red, Blue, Green) configuration):
#   State 1: (0,0,0)   State 2: (0,0,1)   State 3: (0,1,0)   State 4: (0,1,1)
#   State 5: (1,0,0)   State 6: (1,0,1)   State 7: (1,1,0)   State 8: (1,1,1)
# =============================================================================


# ============================================================
# 0) Packages
# ============================================================

library(dplyr)
library(tidyr)
library(stringr)
library(stratEst)


# ============================================================
# 1) Generate dataset (source cleaning script)
# ============================================================

script_dir <- dirname(rstudioapi::getActiveDocumentContext()$path)
source(file.path(script_dir, "DOfile_realexp_3lights (2).R"))

df_strat <- df_final %>%
  rename(
    id     = subject_id,
    period = Guess_Number,
    input  = Light_Config
  ) %>%
  mutate(
    choice    = as.character(Guess),
    treatment = as.character(treatment),
    game      = 1L
  ) %>%
  select(id, period, choice, input, treatment, game) %>%
  arrange(id, treatment, period)

cat("Dataset ready:", nrow(df_strat), "rows |",
    n_distinct(df_strat$id), "subjects |",
    n_distinct(df_strat$treatment), "treatments\n")
cat("Choice values:", paste(sort(unique(df_strat$choice)), collapse = ", "), "\n")
cat("Input values: ", paste(sort(unique(df_strat$input)),  collapse = ", "), "\n")


# ============================================================
# 2) Strategy building blocks
# ============================================================

choice_labels  <- c("0", "1")
input_vals_8   <- c("(0,0,0)", "(0,0,1)", "(0,1,0)", "(0,1,1)",
                    "(1,0,0)", "(1,0,1)", "(1,1,0)", "(1,1,1)")
inputs_strat_8 <- c(NA, input_vals_8)
tr_8state      <- as.numeric(rep(1:8, times = 8))

make_det_strategy_8 <- function(actions) {
  stopifnot(length(actions) == 8, all(actions %in% c(0L, 1L)))
  pc <- as.double(unlist(lapply(actions, function(a) c(1 - a, a))))
  stratEst.strategy(
    choices      = choice_labels,
    inputs       = inputs_strat_8,
    num.states   = 8,
    tr.inputs    = tr_8state,
    prob.choices = pc
  )
}

# Helper: build a named strategy from its integer index j (0-255)
make_by_j <- function(j) {
  make_det_strategy_8(as.integer(intToBits(j))[1:8])
}


# ============================================================
# 3) Build single-variable and two-variable strategies only
# ============================================================

# Integer j → meaningful name.
# Three-variable rules (NOR3, green_alone, blue_alone, blue_green_only,
# red_alone, red_green_only, red_blue_only, AND3, minority, majority,
# ODD_parity, EVEN_parity, NAND3, OR3) are intentionally excluded.

strategy_defs <- c(
  # ── Single-variable ──────────────────────────────────────────────────────
  always_0         =   0L,
  always_1         = 255L,
  green            = 170L,   # sound iff g=1
  not_green        =  85L,   # sound iff g=0
  blue             = 204L,   # sound iff b=1
  not_blue         =  51L,   # sound iff b=0
  red              = 240L,   # sound iff r=1
  not_red          =  15L,   # sound iff r=0

  # ── Two-variable: (red, blue) — green irrelevant ─────────────────────────
  NOR_rb           =   3L,   # r=0 AND b=0
  INHIBIT_br       =  12L,   # b=1 AND r=0
  INHIBIT_rb       =  48L,   # r=1 AND b=0
  EITHER_rb        =  60L,   # r ≠ b   [EITHER treatment]
  NAND_rb          =  63L,   # NOT (r=1 AND b=1)
  AND_rb           = 192L,   # r=1 AND b=1   [AND treatment]
  OR_rb            = 252L,   # r=1 OR b=1    [OR treatment]
  JOINT_rb         = 195L,   # r = b (XNOR)  [JOINT treatment]
  not_INHIBIT_rb   = 207L,   # NOT (r=1 AND b=0)
  not_INHIBIT_br   = 243L,   # NOT (b=1 AND r=0)

  # ── Two-variable: (red, green) — blue irrelevant ─────────────────────────
  NOR_rg           =   5L,   # r=0 AND g=0
  INHIBIT_gr       =  10L,   # g=1 AND r=0
  AND_rg           = 160L,   # r=1 AND g=1
  OR_rg            = 250L,   # r=1 OR g=1
  INHIBIT_rg       =  80L,   # r=1 AND g=0
  EITHER_rg        =  90L,   # r ≠ g
  JOINT_rg         = 165L,   # r = g
  NAND_rg          =  95L,   # NOT (r=1 AND g=1)
  not_INHIBIT_rg   = 175L,   # NOT (r=1 AND g=0)
  not_INHIBIT_gr   = 245L,   # NOT (g=1 AND r=0)

  # ── Two-variable: (blue, green) — red irrelevant ─────────────────────────
  NOR_bg           =  17L,   # b=0 AND g=0
  INHIBIT_gb       =  34L,   # g=1 AND b=0
  AND_bg           = 136L,   # b=1 AND g=1
  OR_bg            = 238L,   # b=1 OR g=1
  INHIBIT_bg       =  68L,   # b=1 AND g=0
  EITHER_bg        = 102L,   # b ≠ g
  JOINT_bg         = 153L,   # b = g
  NAND_bg          = 119L,   # NOT (b=1 AND g=1)
  not_INHIBIT_bg   = 187L,   # NOT (b=1 AND g=0)
  not_INHIBIT_gb   = 221L    # NOT (g=1 AND b=0)
)

stopifnot(!anyDuplicated(strategy_defs))

det_strategies <- setNames(
  lapply(strategy_defs, make_by_j),
  names(strategy_defs)
)

cat("Built", length(det_strategies), "deterministic strategies",
    "(single-variable + two-variable only)\n")


# ============================================================
# 4) Additional strategies: random and always_p
# ============================================================

random <- stratEst.strategy(
  choices      = choice_labels,
  inputs       = inputs_strat_8,
  num.states   = 1,
  tr.inputs    = rep(1, 8),
  prob.choices = c(0.5, 0.5)
)

always_p <- stratEst.strategy(
  choices      = choice_labels,
  inputs       = inputs_strat_8,
  num.states   = 1,
  tr.inputs    = rep(1, 8),
  prob.choices = c(NA_real_, NA_real_)
)


# ============================================================
# 5) Assemble strategy list
# ============================================================

strategies_all <- c(
  det_strategies,
  list(random = random, always_p = always_p)
)

cat("Total strategies:", length(strategies_all), "\n")
stopifnot(!any(sapply(strategies_all, is.null)))
stopifnot(!anyDuplicated(names(strategies_all)))


# ============================================================
# 6) Estimation: loop over all treatments
# ============================================================

all_treatments <- sort(unique(df_strat$treatment))
results_list   <- vector("list", length(all_treatments))
models_list    <- vector("list", length(all_treatments))

for (idx in seq_along(all_treatments)) {
  TREAT <- all_treatments[idx]
  cat("\n========== Treatment:", TREAT,
      "(", idx, "/", length(all_treatments), ") ==========\n")

  df_t <- df_strat %>%
    filter(treatment == TREAT) %>%
    mutate(treatment = as.factor(treatment))

  cat("  Rows:", nrow(df_t), "| Subjects:", n_distinct(df_t$id), "\n")

  model_t <- tryCatch(
    stratEst.model(
      data       = df_t,
      strategies = strategies_all,
      inner.runs = 10,
      outer.runs = 1,
      inner.max  = 10
    ),
    error = function(e) {
      cat("  ERROR:", conditionMessage(e), "\n")
      NULL
    }
  )

  if (!is.null(model_t)) {
    cat("  dim(model_t$shares):", dim(model_t$shares), "\n")
    cat("  sum(model_t$shares):", sum(model_t$shares), "\n")
    shares_vec        <- as.numeric(model_t$shares)
    names(shares_vec) <- names(strategies_all)

    results_list[[idx]] <- data.frame(
      treatment      = TREAT,
      loglike        = model_t$loglike,
      num_par        = model_t$num.par,
      t(shares_vec),
      check.names    = FALSE,
      stringsAsFactors = FALSE
    )
    models_list[[idx]] <- model_t
    cat("  Log-likelihood:", round(model_t$loglike, 3), "\n")
    cat("  Summary:\n")
    print(summary(model_t, plot.shares = FALSE))
  } else {
    results_list[[idx]] <- data.frame(
      treatment = TREAT, loglike = NA_real_, num_par = NA_integer_
    )
  }
}

# Combine results and drop zero-share columns
ESTIMATION_RESULTS <- bind_rows(results_list) %>%
  arrange(treatment)

strat_cols   <- setdiff(names(ESTIMATION_RESULTS), c("treatment", "loglike", "num_par"))
nonzero_cols <- strat_cols[colSums(ESTIMATION_RESULTS[, strat_cols], na.rm = TRUE) > 0]
ESTIMATION_RESULTS <- ESTIMATION_RESULTS[, c("treatment", "loglike", "num_par", nonzero_cols)]

cat("\n=== ESTIMATION_RESULTS (all treatments) ===\n")
print(ESTIMATION_RESULTS %>% mutate(across(where(is.numeric), ~round(., 4))))

output_path <- file.path(script_dir, "Structural_Results_3lights.csv")
write.csv(ESTIMATION_RESULTS, output_path, row.names = FALSE)
cat("Saved to", output_path, "\n")


# ============================================================
# 7) Augment df_final with subject × treatment posteriors
# ============================================================

posterior_list <- vector("list", length(all_treatments))

for (idx in seq_along(all_treatments)) {
  TREAT <- all_treatments[idx]
  mdl   <- models_list[[idx]]
  if (is.null(mdl)) next

  # post.assignment is a subjects × strategies matrix;
  # rownames = id values, colnames = strategy names
  post_mat <- mdl$post.assignment
  if (is.null(post_mat)) next

  post_df            <- as.data.frame(post_mat, row.names = NULL)
  colnames(post_df)  <- paste0("post_", colnames(post_mat))
  post_df$subject_id <- as.integer(sub("^id", "", rownames(post_mat)))
  post_df$treatment  <- TREAT

  posterior_list[[idx]] <- post_df
}

df_posteriors <- bind_rows(posterior_list) %>%
  mutate(treatment = as.character(treatment))

post_cols <- paste0("post_", names(strategies_all))

df_final_post <- df_final %>%
  mutate(treatment = as.character(treatment)) %>%
  left_join(
    df_posteriors %>% select(subject_id, treatment, any_of(post_cols)),
    by = c("subject_id", "treatment")
  )

post_present  <- intersect(post_cols, names(df_final_post))
post_mx       <- as.matrix(select(df_final_post, all_of(post_present)))
best_col      <- max.col(post_mx, ties.method = "first")

df_final_post <- df_final_post %>%
  mutate(
    Rule_used = sub("^post_", "", post_present[best_col]),
    posterior = round(post_mx[cbind(seq_len(nrow(post_mx)), best_col)], 2)
  )

cat("\n=== df_final_post ===\n")
cat("Rows:", nrow(df_final_post), "| Cols:", ncol(df_final_post), "\n")
cat("Posterior columns added:\n")
cat(paste(" ", intersect(post_cols, names(df_final_post)), collapse = "\n"), "\n")
cat("Missing posterior rows:", sum(is.na(df_final_post[[post_cols[1]]])), "\n")


# ============================================================
# 8) Per-treatment share plots (top 5 non-zero strategies)
# ============================================================

for (TREAT in sort(unique(ESTIMATION_RESULTS$treatment))) {
  row    <- ESTIMATION_RESULTS[ESTIMATION_RESULTS$treatment == TREAT, nonzero_cols]
  shares <- as.numeric(row)
  names(shares) <- nonzero_cols
  shares <- head(sort(shares[shares > 0], decreasing = TRUE), 10)

  if (length(shares) == 0) next

  old_mar <- par(mar = c(10, 4, 3, 1))
  barplot(
    shares,
    main   = paste("Estimated shares —", TREAT),
    ylab   = "share",
    ylim   = c(0, 1),
    las    = 2,
    cex.names = 0.75
  )
  par(old_mar)
}